import glob
import json

import numpy as np
import os
import re
from os import listdir
from os.path import join
from scipy import io
import pandas as pd
from torch.utils.data import DataLoader, Dataset
from .process import *

repr_map = {'eventFrame': get_eventFrame,
            'eventAccuFrame': get_eventAccuFrame,
            'timeSurface': get_timeSurface,
            'eventCount': get_eventCount}


# left or right move all event locations randomly
def random_shift_events(events, max_shift=20, resolution=(180, 240)):
    H, W = resolution
    x_shift, y_shift = np.random.randint(-max_shift, max_shift + 1, size=(2,))
    events[:, 1] += x_shift
    events[:, 2] += y_shift

    valid_events = (events[:, 1] >= 0) & (events[:, 1] < W) & (events[:, 2] >= 0) & (events[:, 2] < H)
    events = events[valid_events]

    return events


# flip half of the event images along the x dimension
def random_flip_events_along_x(events, resolution=(180, 240), p=0.5):
    """
    event locations: (t, x, y, p)
    """
    H, W = resolution
    if np.random.random() < p:
        events[:, 1] = W - 1 - events[:, 1]
    return events


class DVS_Lip(Dataset):
    def __init__(self, datafile="../DVS_Lip", phase='train', augmentation=False,
                 repr=['timeSurface', 'eventFrame'], time_num=9,
                 frame_transform=None, label_transform=None):
        self.augmentation = augmentation
        self.datafile = datafile
        self.phase = phase
        self.repr = repr
        self.time_num = time_num

        self.frame_transform = frame_transform
        self.label_transform = label_transform
        self.labels = sorted(os.listdir(os.path.join(datafile, phase)))

        self.file_list = sorted(glob.glob(os.path.join(self.datafile, phase, '*', '*.npy')))

        with open('./data/frame_nums.json', 'r') as f:
            self.frame_nums = json.load(f)

    def __len__(self):
        return len(self.file_list)

    def __getitem__(self, idx):
        """
        returns events and label, loading events from aedat
        :param idx:
        :return: x,y,t,p,  label
        events: (x,y,t,p)  x:[0,346] y:[0,260] p:0/1
        """
        # load timestamps
        word = self.file_list[idx].split('\\')[-2]
        person = self.file_list[idx].split('\\')[-1][:-4]
        # frame_num = self.frame_nums[self.phase][word][int(person)]

        # load events
        try:
            events_input = np.load(self.file_list[idx])
        except:
            print(self.file_list[idx])
        events_input = events_input[np.where(
            (events_input['x'] >= 16) & (events_input['x'] < 112) & (events_input['y'] >= 16) & (
                        events_input['y'] < 112))]
        events_input['x'] -= 16
        events_input['y'] -= 16
        t, x, y, p = events_input['t'], events_input['x'], events_input['y'], events_input['p']
        events_input = np.stack([t, x, y, p], axis=-1)  # [N,4]

        label = self.labels.index(word)

        # normalize the timestamps
        _min = events_input[:, 0].min()
        _max = events_input[:, 0].max()
        events_input[:, 0] = (events_input[:, 0] - _min) / (_max - _min)

        if self.augmentation:
            events_input = random_shift_events(events_input)
            events_input = random_flip_events_along_x(events_input)

        reprs = []
        for repr_name in self.repr:
            # 输入(t,x,y,p), 输出(T,y,x) y:[0,260] x:[0:346]
            repr_array = repr_map[repr_name](events_input[:, 0], events_input[:, 1].astype(np.int32),
                                             events_input[:, 2].astype(np.int32),
                                             events_input[:, 3],
                                             repr_size=(96, 96), time_num=self.time_num)

            # standardization
            mu = np.mean(repr_array)
            sigma = np.std(repr_array)
            repr_array = (repr_array - mu) / sigma

            reprs.append(repr_array)

        reprs = np.array(reprs)  # (C,T,260,346)
        reprs = np.transpose(reprs, (1, 0, 2, 3))  # (T,C,260,346)

        if self.frame_transform is not None:
            reprs = self.frame_transform(reprs)

        if self.label_transform is not None:
            label = self.label_transform(label)

        return reprs, label